import torch
import random

from PIL import Image
from os.path import join
import torch.utils.data as data
import torchvision.transforms as transforms

# this is default dataset dir, this path will be change by the main program
structfile = '/home/a409/users/lihaowei/data/CVUSA/'


class LimitedFoV(torch.nn.Module):
    def __init__(self, fov=360.):
        super().__init__()
        self.fov = fov

    def forward(self, x):
        angle = random.randint(0, 359)
        rotate_index = int(angle / 360. * x.shape[2])
        fov_index = int(self.fov / 360. * x.shape[2])
        if rotate_index > 0:
            img_shift = torch.zeros(x.shape)
            img_shift[:, :, :rotate_index] = x[:, :, -rotate_index:]
            img_shift[:, :, rotate_index:] = x[:, :, :(x.shape[2] - rotate_index)]
        else:
            img_shift = x
        return img_shift[:, :, :fov_index], angle


class Expend(torch.nn.Module):
    def __init__(self, fov=360.):
        super().__init__()
        self.fov = fov

    def forward(self, x):
        fov_index = int(self.fov / 360. * x.shape[2])
        img_shift = x[:, :, : fov_index]
        img_shift = torch.cat((x, img_shift), -1)

        return img_shift


class MaskSky(torch.nn.Module):
    def __init__(self, mask_part=0.5):
        super().__init__()
        self.mask_part = mask_part

    def forward(self, x):
        h = x.shape[1]
        return x[:, int(self.mask_part * h):, :]


def input_transform_grd(fov):
    return transforms.Compose([
        transforms.ToTensor(),
        LimitedFoV(fov),
        # MaskSky(0.5),
        # transforms.Resize([224, 1232]),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])


def input_transform_sat(fov):
    return transforms.Compose([
        transforms.ToTensor(),
        Expend(fov),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])


def get_whole_train_set(mode, fov):
    return WholeDatasetFromStruct(structfile, mode, fov)


class WholeDatasetFromStruct(data.Dataset):
    def __init__(self, structfile, mode, fov):
        super().__init__()

        self.img_root = structfile
        self.data_list = self.img_root + 'splits/'+str(mode)+'.csv'

        # self.input_transform_grd = input_transform_grd(fov)
        self.toTenser = transforms.ToTensor()
        self.limfov = LimitedFoV(fov)
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

        self.input_transform_sat = input_transform_sat(fov)

        self.__cur_id = 0
        self.sat_list = []
        self.grd_list = []
        self.id_list = []
        self.id_idx_list = []
        with open(self.data_list, 'r') as file:
            idx = 0
            for line in file:
                line = line.replace('\"', '')
                data = line.split(',')
                pano_id = (data[0].split('/')[-1]).split('.')[0]
                # satellite filename, streetview filename, pano_id
                item = data[0].replace('bing', 'polar').replace('jpg', 'png')
                grd_inx = data[1]
                self.sat_list.append(item)
                self.grd_list.append(grd_inx)
                self.id_list.append(pano_id)
                self.id_idx_list.append(idx)
                idx += 1
        self.data_size = len(self.id_list)

    def __getitem__(self, index):
        sat_img = Image.open(join(structfile, self.sat_list[index]))
        grd_img = Image.open(join(structfile, self.grd_list[index]))

        sat_img = self.input_transform_sat(sat_img)
        # grd_img = self.input_transform_grd(grd_img)
        grd_img = self.toTenser(grd_img)
        grd_img, angle = self.limfov(grd_img)
        grd_img = self.norm(grd_img)

        return sat_img, grd_img, angle

    def __len__(self):  # 返回列表长度
        return len(self.sat_list)


# test
if __name__ == '__main__':
    pass
    # grd_path = '/home/a409/users/lihaowei/data/CVUSA/streetview/panos/0000001.jpg'
    # grd_img = Image.open(grd_path)
    # trans_grd = input_transform_grd(180)
    # grd_img = trans_grd(grd_img)
    # print(grd_img.shape)
    #
    # sat_path = '/home/a409/users/lihaowei/data/CVUSA/polarmap/19/0000001.png'
    # sat_img = Image.open(sat_path)
    # trans_sat = input_transform_sat(180)
    # sat_img = trans_sat(sat_img)
    # print(sat_img.shape)

    # trans_to_pil = transforms.ToPILImage()
    # img_pil = trans_to_pil(sat_img)
    # img_pil.show()
    # print(5e-4 * (0.95**100))



